package hyman;
import jakarta.xml.bind.JAXBException;
import org.jpmml.evaluator.*;
import org.xml.sax.SAXException;

import javax.xml.parsers.ParserConfigurationException;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;


public class TestMulti {
    public static void main(String[] args) throws JAXBException, IOException, ParserConfigurationException, SAXException {
        // Building a model evaluator from a PMML file
        try (InputStream is = TestMulti.class.getClassLoader().getResourceAsStream("LightGBMAudit_Multiclass.pmml")) {
            if (is == null) {
                throw new FileNotFoundException("Resource LightGBMAudit.pmml not found");
            }

            Evaluator evaluator = new LoadingModelEvaluatorBuilder()
                    .load(is)
                    .build();

            evaluator.verify();
            // Printing input (x1, x2, .., xn) fields
            List<InputField> inputFields = evaluator.getInputFields();
            System.out.println("Input fields: " + inputFields);

            List<TargetField> targetFields = evaluator.getTargetFields();
            System.out.println("Target field(s): " + targetFields);

            List<OutputField> outputFields = evaluator.getOutputFields();
            System.out.println("Output fields: " + outputFields);

            // 构造并评估一条记录

            // 构造一条测试记录
            Map<String, Object> inputRecord = createRecord(
                    38, "Private", "College", "Married", "Farming", 22099.0, "Male", 40
            );

            Map<String, ?> results = evaluator.evaluate(inputRecord);
            results = EvaluatorUtil.decodeAll(results);

            Object predictedLabel = results.get("None");
            System.out.println("Predicted Label: " + predictedLabel);

// 提取各分类的概率
            Map<String, Double> probabilities = new LinkedHashMap<>();
            for (Map.Entry<String, ?> entry : results.entrySet()) {
                if (entry.getKey().startsWith("probability(")) {
                    String className = entry.getKey()
                            .replace("probability(", "")
                            .replace(")", "");
                    Double value = Double.parseDouble(entry.getValue().toString());
                    probabilities.put(className, value);
                }
            }

            System.out.println("Probabilities:");
            for (Map.Entry<String, Double> entry : probabilities.entrySet()) {
                System.out.printf("  %s: %.6f%n", entry.getKey(), entry.getValue());
            }

            evaluator = null;

        }


    }

    static Map<String, Object> createRecord(Object age, String employment, String education,
                                            String marital, String occupation, Object income,
                                            String gender, Object hours) {
        Map<String, Object> record = new LinkedHashMap<>();
        record.put("Age", age);         // 支持 null 表示缺失
        record.put("Employment", employment);
        record.put("Education", education);
        record.put("Marital", marital);
        record.put("Occupation", occupation);
        record.put("Income", income);   // 支持数值类型（Double/Integer）或 null
        record.put("Gender", gender);
        record.put("Hours", hours);     // 支持数值类型
        return record;
    }
}
